from cmath import sqrt
import os
import numpy as np
import time
import matplotlib.pyplot as plt
import torch
import torchvision


class Plot:
    """
    the summary figure plot class for ploting the inner or the final experiment result.
    """
    def __init__(self):

        self.root_path = './img'
        self.path = os.path.join(self.root_path, 'Experiment')
        if os.path.exists(self.path) == False:
            os.makedirs(self.path)
        

    def plot_archive(self, archive):

        plt.title('Achive size : {}'.format(archive['size']))

        arr1 = []
        arr2 = []

        for particle in archive['set']:

            arr1.append(particle.fitness[0])
            arr2.append(particle.fitness[1])
        

        plt.scatter(arr1, arr2)
        plt.xlabel('f1')
        plt.ylabel('f2')

        plt.savefig(os.path.join(self.path, 'test1_{}.jpg'.format(time.mktime(time.gmtime()))))
        plt.close()
        print('Figure save done')
    
    def plot_image(self, image, name):

        def imshow(img):
            np_img = img.cpu().numpy()
            plt.imshow(np.transpose(np_img, (1, 2, 0)))
        imshow(torchvision.utils.make_grid(image))
        plt.savefig(os.path.join(self.path, 'figure_{}.jpg'.format(name)))
        plt.close()
        print('Figure save done')

def dict_show(dict):

    print('=='*50)
    for key, value in dict.items():
        print(key, end=' :')
        print(value)
    print('=='*50)
    

class DemoTest:

    def __init__(self, func_name):

        self.func_name = func_name
        self.param = self.get_param()
    
    def get_param(self, ):

        dicter = {
            'ZDT1': {
                'DsDim': 30,
                'OsDim': 2,
            }
        } 


        return dicter[self.func_name]


    def ZDT1(self, x):
        
        if x.shape[0] != self.param['DsDim']:
            raise ValueError
        
        f1 = x[0]
        g = 1 + 9 * (x[1:].sum() / (x.shape[0] - 1))
        h = 1 - sqrt(f1 / g)
        f2 = g * h

        return np.array([f1, f2])

